import uvicorn
import time
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
# 导入我们的ASR模块和路由器
from dolphin_asr import dolphin_router, load_model_once
# 导入WebSocket路由器
from audio_websocket import router as websocket_router
from contextlib import asynccontextmanager
# 仅导入需要的部分，不再导入可能导致循环依赖的内容
from punc_restore import punc_router
from config import *

# 全局变量 - 服务启动时加载模型
global_dolphin_model = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时加载模型
    global global_dolphin_model
    
    try:
        print("服务启动，预加载ASR模型...")
        global_dolphin_model = load_model_once(device="cuda")
        print("ASR模型加载完成")
        
        # 导入并设置punc_restore中的全局变量
        if global_load_punc_model:
            import punc_restore
            
            # 加载标点恢复模型
            print("开始加载标点恢复模型...")
            punc_model = punc_restore.load_punc_model()
            
            # 设置punc_restore中的全局变量
            punc_restore.global_punc_model = punc_model
            
            # 验证模型是否正确加载
            if punc_model is None:
                print("警告: 标点恢复模型加载失败，将在首次使用时重试")
            else:
                print("标点恢复模型加载成功")
                
                # 预热调用，避免首次请求延迟过高
                print("执行标点恢复模型预热...")
                warm_up_start = time.time()
                
                # 使用配置文件中的预热文本列表
                for i, text in enumerate(punc_preload_texts, 1):
                    print(f"预热调用 {i}/{len(punc_preload_texts)}...")
                    result = punc_restore.text_punc_restore(text, model=punc_model)
                
                warm_up_end = time.time()
                print(f"标点恢复模型预热完成，总耗时: {warm_up_end - warm_up_start:.4f} 秒")
                print(f"预热后缓存中的文本数量: {len(punc_restore.result_cache)}")
        
        print("所有模型加载完成，服务就绪")
    except Exception as e:
        print(f"模型加载过程中出现错误: {str(e)}")
        print("服务将继续启动，但某些功能可能不可用")
        # 不要在这里重新抛出异常，让服务继续启动
    
    yield
    # 关闭时的清理代码可以放在这里

# 创建FastAPI应用
app = FastAPI(
    title="Dolphin ASR API",
    description="语音识别API基于Dolphin ASR模型",
    version="1.0.0",
    lifespan=lifespan
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源，生产环境中应该设置为特定的域名
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)

# 包含ASR路由器
app.include_router(dolphin_router, prefix="/asr-api")

# 包含WebSocket路由器
app.include_router(websocket_router, prefix="/ws")

# 包含标点恢复路由器
app.include_router(punc_router, prefix="/punc-api")

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")

# 添加根路径重定向到主页
@app.get("/")
async def root():
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/static/index.html")

# 添加API文档路径
@app.get("/api-docs")
async def api_docs():
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/docs")

if __name__ == "__main__":
    # 启动FastAPI服务
    uvicorn.run(app, host="0.0.0.0", port=8000) 